{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 2: Analyzing LAMMPS thermodynamic data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Author: [Richard Berger](mailto:richard.berger@outlook.com)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial assumes you've completed the [first example](simple.ipynb) and understand the basics of running LAMMPS through Python. In this tutorial we will build on top of that example and look at how to extract thermodynamic data produced by LAMMPS into Python and visualize it. Let's first start by recreating our simple melt example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext lammps.ipython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lammps import lammps\n",
    "L = lammps()\n",
    "L.cmd.auto_flush = True\n",
    "\n",
    "def init_melt_system(L):\n",
    "    # 3d Lennard-Jones melt\n",
    "    L.cmd.clear()\n",
    "    L.cmd.units(\"lj\")\n",
    "    L.cmd.atom_style(\"atomic\")\n",
    "    \n",
    "    L.cmd.lattice(\"fcc\", 0.8442)\n",
    "    L.cmd.region(\"box\", \"block\", 0, 4, 0, 4, 0, 4)\n",
    "    L.cmd.create_box(1, \"box\")\n",
    "    L.cmd.create_atoms(1, \"box\")\n",
    "    L.cmd.mass(1, 1.0)\n",
    "    \n",
    "    L.cmd.velocity(\"all\", \"create\", 1.44, 87287, \"loop geom\")\n",
    "    \n",
    "    L.cmd.pair_style(\"lj/cut\", 2.5)\n",
    "    L.cmd.pair_coeff(1, 1, 1.0, 1.0, 2.5)\n",
    "    \n",
    "    L.cmd.neighbor(0.3, \"bin\")\n",
    "    L.cmd.neigh_modify(\"delay\", 0, \"every\", 20, \"check no\")\n",
    "    \n",
    "    L.cmd.fix(\"1\", \"all\", \"nve\")\n",
    "    \n",
    "    L.cmd.thermo(50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we take advantage of the fact that we can write regular Python functions to organize our LAMMPS simulation. This allows us to clear and initialize a new system by calling the `init_melt_system()` function. With this we can now go ahead an run this simulation for 100 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_melt_system(L)\n",
    "L.cmd.run(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting thermodynamic data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at the above output we see that LAMMPS prints out thermodynamic data for steps 0, 50 and 100.\n",
    "\n",
    "```\n",
    "   Step          Temp          E_pair         E_mol          TotEng         Press     \n",
    "         0   1.44          -6.7733681      0             -4.6218056     -5.0244179    \n",
    "        50   0.70303849    -5.6796164      0             -4.629178       0.50453907   \n",
    "       100   0.72628044    -5.7150774      0             -4.6299123      0.29765862\n",
    "```\n",
    "\n",
    "We could parse the text output and extract the necessary information, but this has proven to be error-prone and clunky, especially in cases where other output gets interleaved with thermo output lines. Instead, we can make use of the Python integration within LAMMPS to execute arbitrary Python code during time steps using `fix python/invoke`. We can extract the thermodynamic data directly using the LAMMPS Python interface and process it in any way we want.\n",
    "\n",
    "For this we first define the data structure we want to use to store the data. For each column of the thermodynamic data we want to store a list of values for each time step. Let's use a Python `dict` with the following structure:\n",
    "\n",
    "```python\n",
    "{'Step': [0, 50, 100, ...], 'Temp': [...], 'E_pair': [...], 'E_mol': [...], 'TotEng': [...], 'Press': [...]}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, let's define an empty `dict` and call it `current_run`. As the simulation progresses, we append new data into this dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_run = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's define a function that should be executed every time step a thermodynamic output line would be written. This function takes a `lammps` class instance and through it can access LAMMPS state and data. We can use the [`last_thermo()`](https://docs.lammps.org/Python_module.html#lammps.lammps.last_thermo) function of the `lammps` class to get the latest thermodynamic data as a dictionary. This data is all we need to populate our `current_run` data structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def append_thermo_data(lmp):\n",
    "  for k, v in lmp.last_thermo().items():\n",
    "    current_run.setdefault(k, []).append(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With these two pieces in place, it is now time to tell LAMMPS about how we want to call this function.\n",
    "\n",
    "First, let's suppress any LAMMPS output via `%%capture_lammps_output` and reinitialize our system with `init_melt_system()` so our system is back in its initial state and the time step is back to 0.\n",
    "\n",
    "Next, we add a new fix `python/invoke` that should execute every 50 time steps, the same as our `thermo 50` command above. At the end of every 50 time steps (including the first one), it should call the `append_thermo_data` function we just defined. Notice we can just pass the function as parameter. Finally, we tell LAMMPS to run for 250 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture_lammps_output\n",
    "init_melt_system(L)\n",
    "L.cmd.fix(\"myfix\", \"all\", \"python/invoke\", 50, \"end_of_step\", append_thermo_data)\n",
    "L.cmd.run(250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect our `current_run` dictionary after the run has completed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the time steps 0, 50, 100, 150, and 200 were added to dictionary. However, the last time step 250 is still missing. For this we need to manually add a final call to our `append_thermo_data()` helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "append_thermo_data(L)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this our `current_run` dictionary now has all the data of the completed run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting thermodynamic data with matplotlib\n",
    "\n",
    "Now that we have our data available as Python variables, we can easily use other libraries for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.xlabel('time step')\n",
    "plt.ylabel('Total Energy')\n",
    "plt.plot(current_run['Step'], current_run['TotEng'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Pandas library\n",
    "\n",
    "Since we can call any Python code from LAMMPS, the above example can also be rewritten using the Pandas library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture_lammps_output\n",
    "import pandas as pd\n",
    "\n",
    "current_run = pd.DataFrame()\n",
    "\n",
    "def append_thermo_data(lmp):\n",
    "    global current_run\n",
    "    current_time_step = pd.DataFrame.from_records([lmp.last_thermo()])\n",
    "    current_run = pd.concat([current_run, current_time_step], ignore_index=True)\n",
    "\n",
    "init_melt_system(L)\n",
    "L.cmd.fix(\"myfix\", \"all\", \"python/invoke\", 50, \"end_of_step\", append_thermo_data)\n",
    "L.cmd.run(250)\n",
    "append_thermo_data(L)\n",
    "current_run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_run.plot(x='Step', y='TotEng')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "The Python interface gives you a powerful way of invoking and extracting simulation data while the simulation is running. Next we'll look at how to extract information about the atoms in your system.\n",
    "\n",
    "<div style=\"text-align:right\"><a href=\"atoms.ipynb\">Next</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
